# encoding:utf8
import logging
import os
import pandas as pd

logger = logging.Logger


class CSVDatafeed:
    def __init__(self):
        self.code_dfs = {}

    def add_data(self, code, csv_file):
        if not os.path.exists(csv_file):
            logging.error('{}csv文件不存在！'.format(code))
            return

        df = pd.read_csv(csv_file, index_col=0)
        if len(df) == 0:
            logging.error('{}没有数据！'.format(code))
            return

        for col in ['date']:
            if col not in df.columns:
                logging.error('{}字段{}不存在！'.format(code, col))
                return

        df['code'] = code
        df['date'] = df['date'].apply(lambda x: str(x))
        df.sort_values(by='date', inplace=True)

        df['rate'] = df['close'].pct_change()
        df['equity'] = (df['rate']+1).cumprod()

        self.code_dfs[code] = df

    def get_df(self, instrument):
        if instrument in self.code_dfs.keys():
            df = self.code_dfs[instrument]
            df.index = df['date']
            return df
        else:
            logging.info('{}未加载'.format(instrument))
            return None

    def get_all_df(self):
        df_all = pd.concat(self.code_dfs.values(), axis=0)
        df_all.dropna(inplace=True)
        df_all.index = df_all['date']
        df_all.sort_index(inplace=True)

        return df_all


if __name__ == '__main__':
    feed = CSVDatafeed()
    feed.add_data('000300.OF', '../datas/000300.SH.csv')
    feed.add_data('000905.OF', '../datas/000300.SH.csv')
    df = feed.get_all_df()
    print(df)
